package com.xjtu.mybatisplus.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.completion.chat.*;
import com.theokanning.openai.service.OpenAiService;
import com.xjtu.mybatisplus.model.dto.ChatDTO;
import com.xjtu.mybatisplus.model.dto.CreateChatDTO;
import io.reactivex.Flowable;
import okhttp3.OkHttpClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import retrofit2.Retrofit;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static com.theokanning.openai.service.OpenAiService.*;

@Service
public class ChatService {
    OpenAiService openAiService;

    Map<Integer, List<ChatMessage>> chatMap;

    public ChatService() {
        String host = "https://openai.jsong.site";
        Duration timeout = Duration.ofSeconds(10);
        ObjectMapper mapper = defaultObjectMapper();
        String token = "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";
        OkHttpClient client = defaultClient(token, timeout)
                .newBuilder()
                .build();
        Retrofit retrofit = defaultRetrofit(client, mapper).newBuilder()
                .baseUrl(host).build();

        OpenAiApi api = retrofit.create(OpenAiApi.class);
        openAiService = new OpenAiService(api);

        chatMap = new HashMap<>();
    }

    public SseEmitter streamChat(ChatDTO chatDTO) {
        final SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
        List<ChatMessage> messages = chatMap.get(chatDTO.getChatId());
        messages.add(new ChatMessage(ChatMessageRole.USER.value(), chatDTO.getMessage()));
        ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
                .model("gpt-3.5-turbo-1106")
                .messages(messages)
                .maxTokens(2560)
                .build();
        Flowable<ChatCompletionChunk> flowAble = openAiService.streamChatCompletion(completionRequest);
        StringBuilder sb = new StringBuilder();
        openAiService.mapStreamToAccumulator(flowAble)
                .doOnNext(chatMessageAccumulator -> {
                    String content = chatMessageAccumulator.getMessageChunk().getContent();
                    if (content != null) {
                        try {
                            emitter.send(content.getBytes(StandardCharsets.UTF_8));
                            System.out.print(content);
                            sb.append(content);

                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                })
                .doOnComplete(() -> {
                    emitter.complete();
                    messages.add(new ChatMessage(ChatMessageRole.SYSTEM.value(), sb.toString()));
                })
                .subscribe();
        return emitter;
    }

    public String chat(ChatDTO chatDTO) {
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage(ChatMessageRole.USER.value(), chatDTO.getMessage()));
        ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
                .model("gpt-3.5-turbo-1106")
                .messages(messages)
                .maxTokens(4096)
                .build();
        ChatCompletionResult completion = openAiService.createChatCompletion(completionRequest);
        return completion.getChoices().get(0).getMessage().getContent();
    }

    public String summary(String content) {
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage(ChatMessageRole.SYSTEM.value(), "你是一个报纸的编辑，用户给你传递一篇文章后你需要生成这篇文章的一个封面介绍语。你的回复仅应当包含一个介绍语句，长度在200字左右。回复应当是幽默生动有趣的。"));
        messages.add(new ChatMessage(ChatMessageRole.USER.value(), content));
        ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
                .model("gpt-3.5-turbo-1106")
                .messages(messages)
                .maxTokens(4096)
                .build();
        ChatCompletionResult completion = openAiService.createChatCompletion(completionRequest);
        return completion.getChoices().get(0).getMessage().getContent();
    }

    public static void main(String[] args) {
        ChatService chatService = new ChatService();
    }

    @Autowired
    private ArticleService articleService;

    public Integer createChat(CreateChatDTO chatDTO) {
        Integer chatId = chatMap.size();
        // TODO: 这里会有并发问题，需要用原子ID
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage(ChatMessageRole.SYSTEM.value(),
                "你接下来会收到一篇文章，随后你需要根据这篇文章的内容对用户的提问进行回复，回复应当是中文，纯文本。"));
        String content = articleService.getById(chatDTO.getArticleId()).getContent();
        messages.add(new ChatMessage(ChatMessageRole.USER.value(), content));
        chatMap.put(chatId, messages);
        return chatId;
    }
}
